package com.walking.slf4jmdc.threadexecutor;

import org.slf4j.MDC;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;

import java.util.Map;
import java.util.UUID;
import java.util.concurrent.Future;

/**
 * @author walking
 */
public class MDCThreadExecutor extends ThreadPoolTaskExecutor {
    private Map<String, String> threadContext;
    private Boolean useThreadContext;

    public MDCThreadExecutor() {
        super();
    }

    public MDCThreadExecutor(Map<String, String> threadContext) {
        this.threadContext = threadContext;
        this.useThreadContext = threadContext != null;
    }

    public Map<String, String> getThreadContext() {
        return useThreadContext ? (threadContext == null ? MDC.getCopyOfContextMap() : threadContext) : null;
    }

    @Override
    public void execute(Runnable task) {
        super.execute(getMDCTask(task, getThreadContext()));
    }
    @Override
    public Future<?> submit(Runnable task) {
        return super.submit(getMDCTask(task,getThreadContext()));
    }

    private Runnable getMDCTask(Runnable task, Map<String, String> threadContext) {
        return () -> {
            if (useThreadContext) {
                MDC.setContextMap(threadContext);
                MDC.put("tranceID", UUID.randomUUID().toString().replaceAll("-",""));
            }else{
                MDC.setContextMap(null);
            }
            try {
                task.run();
            } finally {
                MDC.remove("tranceID");
                MDC.clear();
            }
        };
    }


}
